import os
from glob import glob

import SimpleITK as sitk
import nibabel
import numpy
import numpy as np

import config
from utils.LogUtil import my_logger


def get_all_data(data_dir_list, patch_shape, crop_method, if_mask=True):
    """
    Try to read original image data, no data augmentation, no unsqueeze
    :param data_dir_list: The list of data's path
    :param patch_shape:  patch shape
    :param crop_method: 1&2 mean crop the data to many patches,  3 means crop the data to 1 patch
    When the crop_method is 3, parameter patch_shape will be invalid
    :param if_mask:
    :return:
    """
    my_logger.info("start to read data...")
    img_list = []
    label_list = []

    for i, sample_dir in enumerate(data_dir_list):
        if if_mask:
            tmp_data_list, tmp_label_list = get_patch_with_mask(sample_dir,
                                                                patch_shape,
                                                                crop_method)
        else:
            tmp_data_list, tmp_label_list = get_patch_without_mask(sample_dir,
                                                                   patch_shape,
                                                                   crop_method)
        my_logger.info(
            str(len(tmp_data_list)) + " patches of " + os.path.basename(sample_dir) + " has been added to list")
        img_list += tmp_data_list
        label_list += tmp_label_list
    preprocess(img_list)
    my_logger.info("The total number of data is " + str(len(img_list)))
    return img_list, label_list


def get_patch(sample_dir, target_shape, if_crop):
    """
    This function is obsolete, because it can change the code depending on whether the data have mask
    :param sample_dir:
    :param target_shape:
    :param if_crop:
    :return:
    """
    train_img, label = read_nii_file(sample_dir)  # the result's shape is (n_channel, X, Y, X)
    tmp = np.max(label)
    if np.max(label) == 0:
        my_logger.warn("discard a patch")
        return
    if if_crop == 0:
        train_result_list = [train_img]
        label_result_list = [label]
        return train_result_list, label_result_list
    # The shape of train_img and label must be z,x,y (z is the number os slices)
    train_result_list = []
    label_result_list = []
    original_shape = train_img.shape
    z_num = original_shape[1] // target_shape[0]
    x_num = original_shape[2] // target_shape[1]
    y_num = original_shape[3] // target_shape[2]
    if if_crop == 1:  # random crop
        for i in range(z_num * x_num * y_num):
            random_z_index = np.random.randint(0, high=original_shape[1] - target_shape[0] + 1)
            random_x_index = np.random.randint(0, high=original_shape[2] - target_shape[1] + 1)
            random_y_index = np.random.randint(0, high=original_shape[3] - target_shape[2] + 1)
            train_result_list.append(train_img[:, random_z_index:random_z_index + target_shape[0],
                                     random_x_index:random_x_index + target_shape[1],
                                     random_y_index:random_y_index + target_shape[2]])
            label_result_list.append(label[:, random_z_index:random_z_index + target_shape[0],
                                     random_x_index:random_x_index + target_shape[1],
                                     random_y_index:random_y_index + target_shape[2]])
    elif if_crop == 2:  # direct crop
        for i in range(z_num):
            for j in range(x_num):
                for k in range(y_num):
                    train_result_list.append(train_img[:, i * target_shape[0]:(i + 1) * target_shape[0],
                                             j * target_shape[1]:(j + 1) * target_shape[1],
                                             k * target_shape[2]:(k + 1) * target_shape[2]])
                    label_result_list.append(label[:, i * target_shape[0]:(i + 1) * target_shape[0],
                                             j * target_shape[1]:(j + 1) * target_shape[1],
                                             k * target_shape[2]:(k + 1) * target_shape[2]])

    return train_result_list, label_result_list


def get_patch_without_mask(sample_dir, patch_shape, crop_method, if_label=True):
    train_data, label = read_nii_file(sample_dir, if_label)
    if crop_method == 4:
        return [train_data], [label]
    original_shape = train_data.shape
    z_num = int(np.ceil(original_shape[0] / patch_shape[0]))
    x_num = int(np.ceil(original_shape[1] / patch_shape[1]))
    y_num = int(np.ceil(original_shape[2] / patch_shape[2]))
    if original_shape[0] < patch_shape[0] or original_shape[1] < patch_shape[1] or original_shape[2] < patch_shape[2]:
        my_logger.error("The patch size is too big!")
        raise Exception
    elif crop_method == 3:  # crop to 1 patch
        original_shape = train_data.shape
        new_roi = get_appropriate_shape(original_shape)
        train_data = train_data[new_roi[0][0]:new_roi[0][1], new_roi[1][0]:new_roi[1][1], new_roi[2][0]:new_roi[2][1]]
        label = label[new_roi[0][0]:new_roi[0][1], new_roi[1][0]:new_roi[1][1], new_roi[2][0]:new_roi[2][1]]
        return [train_data], [label]
    elif crop_method == 2:  # direct crop
        train_result_list = []
        label_result_list = []
        for i in range(z_num):
            z_start = i * patch_shape[0]
            z_end = (i + 1) * patch_shape[0]
            if z_end > original_shape[0]:  # the last one
                z_end = original_shape[0]
                z_start = original_shape[0] - patch_shape[0]
            for j in range(x_num):
                x_start = j * patch_shape[1]
                x_end = (j + 1) * patch_shape[1]
                if x_end > original_shape[1]:  # the last one
                    x_end = original_shape[1]
                    x_start = original_shape[1] - patch_shape[1]
                for k in range(y_num):
                    y_start = k * patch_shape[2]
                    y_end = (k + 1) * patch_shape[2]
                    if y_end > original_shape[2]:  # the last one
                        y_end = original_shape[2]
                        y_start = original_shape[2] - patch_shape[2]
                    print(z_start, z_end, x_start, x_end, y_start, y_end)
                    train_result_list.append(
                        train_data[z_start:z_end, x_start:x_end, y_start:y_end])
                    if label:
                        label_result_list.append(
                            train_data[z_start:z_end, x_start:x_end, y_start:y_end])
        return train_result_list, label_result_list
    elif crop_method == 1:  # random crop
        train_result_list = []
        label_result_list = []
        for i in range(int(z_num * x_num * y_num)):  # 2 means that cropping more train data
            random_z_index = np.random.randint(0, high=original_shape[0] - patch_shape[0] + 1)
            random_x_index = np.random.randint(0, high=original_shape[1] - patch_shape[1] + 1)
            random_y_index = np.random.randint(0, high=original_shape[2] - patch_shape[2] + 1)
            train_result_list.append(train_data[random_z_index:random_z_index + patch_shape[0],
                                     random_x_index:random_x_index + patch_shape[1],
                                     random_y_index:random_y_index + patch_shape[2]])
            label_result_list.append(label[random_z_index:random_z_index + patch_shape[0],
                                     random_x_index:random_x_index + patch_shape[1],
                                     random_y_index:random_y_index + patch_shape[2]])
        return train_result_list, label_result_list
    else:
        my_logger.error("Wrong crop method code!")
        raise Exception


def get_patch_with_mask(sample_dir, target_shape, crop_method):
    train_result_list = []
    label_result_list = []
    train_patched_file_list = glob(os.path.join(sample_dir, "image", "*patch.npy"))
    label_patched_file_list = glob(os.path.join(sample_dir, "label", "*patch.npy"))

    if len(train_patched_file_list) > 0:  # This sample has patched
        assert (len(train_result_list) == len(label_result_list))
        list.sort(train_patched_file_list)
        list.sort(label_patched_file_list)
        for train_patched_file, label_patched_file in zip(train_patched_file_list, label_patched_file_list):
            train_patch_tmp = np.load(train_patched_file)
            label_patch_tmp = np.load(label_patched_file)
            # label_patch_tmp = label_patch_tmp.astype(numpy.long)  # one hot need the type is long
            if np.max(train_patch_tmp) - np.min(train_patch_tmp) > 0 and np.max(label_patch_tmp) - np.min(
                    label_patch_tmp) > 0:
                # train_patch_tmp = preprocess(train_patch_tmp)
                train_result_list.append(train_patch_tmp)
                label_result_list.append(label_patch_tmp)
                # my_logger.info("reading train patch from " + train_patched_file)
                # my_logger.info("reading label patch from " + label_patched_file)
        assert (len(train_result_list) == len(label_result_list))
        return train_result_list, label_result_list

    # save_root_path = r"D:\gs\code\parse2022-draft\resource\train_patch_sample"
    # the result's shape is (n_channel, X, Y, X)
    train_img, label, train_I_affine, label_I_affine, _ = read_nii_file_with_mask(sample_dir)
    if np.max(label) == 0:
        my_logger.warn("discard a patch")
        return
    if crop_method == 0:
        train_result_list = [train_img]
        label_result_list = [label]
        return train_result_list, label_result_list
    # The shape of train_img and label must be z,x,y (z is the number os slices)

    original_shape = train_img.shape
    z_num = original_shape[0] // target_shape[0]
    x_num = original_shape[1] // target_shape[1]
    y_num = original_shape[2] // target_shape[2]
    if z_num == 0 or x_num == 0 or y_num == 0:
        my_logger.info("crop failed : " + sample_dir)
        return [], []
    if crop_method == 1:  # random crop
        abandon_number = 0
        for i in range(z_num * x_num * y_num):
            random_z_index = np.random.randint(0, high=original_shape[0] - target_shape[0] + 1)
            random_x_index = np.random.randint(0, high=original_shape[1] - target_shape[1] + 1)
            random_y_index = np.random.randint(0, high=original_shape[2] - target_shape[2] + 1)
            # if the number of artery voxel is too small
            if np.sum(label[:, random_z_index:random_z_index + target_shape[0],
                      random_x_index:random_x_index + target_shape[1],
                      random_y_index:random_y_index + target_shape[2]]) < 0.01 * target_shape[0] * target_shape[1] * \
                    target_shape[2]:
                abandon_number += 1
                continue
            train_result_list.append(train_img[:, random_z_index:random_z_index + target_shape[0],
                                     random_x_index:random_x_index + target_shape[1],
                                     random_y_index:random_y_index + target_shape[2]])
            label_result_list.append(label[:, random_z_index:random_z_index + target_shape[0],
                                     random_x_index:random_x_index + target_shape[1],
                                     random_y_index:random_y_index + target_shape[2]])

            # temp_train_img = np.zeros_like(train_img)
            # temp_train_img[:,
            # random_z_index:random_z_index + target_shape[0],
            # random_x_index:random_x_index + target_shape[1],
            # random_y_index:random_y_index + target_shape[2]] = train_img[:,
            #                                                    random_z_index:random_z_index + target_shape[0],
            #                                                    random_x_index:random_x_index + target_shape[1],
            #                                                    random_y_index:random_y_index + target_shape[2]]
            # temp_train_img_save_path = os.path.join(save_root_path, os.path.basename(sample_dir) + "_train_" + str(i) + ".nii.gz")
            # nibabel.Nifti1Image(temp_train_img[0], train_I_affine).to_filename(temp_train_img_save_path)
            # my_logger.info("save train patch to " + temp_train_img_save_path)
            #
            # temp_label_img = np.zeros_like(label)
            # temp_label_img[:,
            # random_z_index:random_z_index + target_shape[0],
            # random_x_index:random_x_index + target_shape[1],
            # random_y_index:random_y_index + target_shape[2]] = label[:,
            #                                                    random_z_index:random_z_index + target_shape[0],
            #                                                    random_x_index:random_x_index + target_shape[1],
            #                                                    random_y_index:random_y_index + target_shape[2]]
            # temp_label_img_save_path = os.path.join(save_root_path,
            #                                         os.path.basename(sample_dir) + "_label_" + str(i) + ".nii.gz")
            # nibabel.Nifti1Image(temp_label_img[0], label_I_affine).to_filename(temp_label_img_save_path)
            # my_logger.info("save label patch to "+temp_label_img_save_path)

        my_logger.info("The abandon number is " + str(abandon_number))
        my_logger.info(
            "The abandon ratio is " + str(abandon_number / (z_num * x_num * y_num)))

    elif crop_method == 2:  # direct crop
        for i in range(z_num):
            for j in range(x_num):
                for k in range(y_num):
                    train_result_list.append(train_img[:, i * target_shape[0]:(i + 1) * target_shape[0],
                                             j * target_shape[1]:(j + 1) * target_shape[1],
                                             k * target_shape[2]:(k + 1) * target_shape[2]])
                    label_result_list.append(label[:, i * target_shape[0]:(i + 1) * target_shape[0],
                                             j * target_shape[1]:(j + 1) * target_shape[1],
                                             k * target_shape[2]:(k + 1) * target_shape[2]])

    return train_result_list, label_result_list


def read_nii_file_detail(data_path):
    sample_name = os.path.basename(data_path)
    train_data_path = os.path.join(data_path, "image", sample_name + ".nii.gz")
    label_data_path = os.path.join(data_path, "label", sample_name + ".nii.gz")
    train_I = nibabel.load(train_data_path)
    train_I_affine = train_I.affine
    train_data = train_I.get_data()
    label_I = nibabel.load(label_data_path)
    label_I_affine = label_I.affine
    label_data = label_I.get_data()
    my_logger.info("read data detail from " + data_path)
    return train_I_affine, train_data, label_I_affine, label_data


def read_nii_file(data_path, if_label=True):
    if "mri" and "Preprocessing" in data_path:
        sample_name = os.path.basename(data_path)
        train_data_path = os.path.join(data_path, "MRA.nii.gz")
        train_nii_img = nibabel.load(train_data_path)
        train_img = train_nii_img.get_data().transpose((2, 1, 0))
        # train_img = preprocess(train_img)
        mask_data_path = os.path.join(data_path, "..", "..", "TrainMask", "train" + sample_name[-2:] + ".nii.gz")
        mask_nii_img = nibabel.load(mask_data_path)
        mask = mask_nii_img.get_data().astype(np.int).transpose((2, 1, 0))
        train_img = train_img * mask
        if if_label:
            label_path = os.path.join(data_path, "label.nii.gz")
            label_nii_img = nibabel.load(label_path)
            label = label_nii_img.get_data().transpose((2, 1, 0))
            label = label.astype(np.int)
        else:
            label = None
        my_logger.info("read data from " + data_path)
        return train_img, label

    if "mri" in data_path:
        sample_name = os.path.basename(data_path)  # 128 * 448 * 448
        train_data_path = os.path.join(data_path, "image", sample_name + ".mhd")
        train_img = sitk.ReadImage(train_data_path)
        train_img = sitk.GetArrayFromImage(train_img)
        train_img = preprocess(train_img)
        mask_data_path = os.path.join(data_path, "..", "..", "TrainMask", sample_name + ".nii.gz")
        mask_nii_img = nibabel.load(mask_data_path)
        mask = mask_nii_img.get_data().astype(np.int).transpose((2, 1, 0))
        train_img = train_img * mask
        if if_label:
            label_data_path = os.path.join(data_path, "label", sample_name + ".mhd")
            label = sitk.ReadImage(label_data_path)
            label = sitk.GetArrayFromImage(label)
            label = label.astype(np.int)
        else:
            label = None
        my_logger.info("read data from " + data_path)
        return train_img, label
    else:
        sample_name = os.path.basename(data_path)
        train_data_path = os.path.join(data_path, "image", sample_name + ".nii.gz")
        label_data_path = os.path.join(data_path, "label", sample_name + ".nii.gz")
        train_data = nibabel.load(train_data_path).get_data()
        label_data = nibabel.load(label_data_path).get_data()
        my_logger.info("read data from " + data_path)
        return train_data, label_data


def read_nii_file_with_mask(data_path):
    sample_name = os.path.basename(data_path)
    train_I_affine, train_data, label_I_affine, label_data = read_nii_file_detail(data_path)
    mask_data_path = os.path.join(data_path, "image", sample_name + "_lung" + ".nii.gz")

    mask_data = nibabel.load(mask_data_path).get_data()
    mask_test = np.where(mask_data)
    Box_boundary = [np.min(mask_test[0]), np.max(mask_test[0]) + 1, np.min(mask_test[1]), np.max(mask_test[1]) + 1,
                    np.min(mask_test[2]), np.max(mask_test[2]) + 1]
    # train_data = (train_data - np.min(train_data)) / (np.max(train_data) - np.min(train_data)) * 255

    # get the box boundary
    Box = np.zeros(shape=train_data.shape)
    Box[Box_boundary[0]:Box_boundary[1], Box_boundary[2]:Box_boundary[3],
    Box_boundary[4]:Box_boundary[5]] = 1
    train_data = np.where(Box == 1, train_data, np.min(train_data))
    label_data = np.where(Box == 1, label_data, np.min(label_data))

    # nibabel.Nifti1Image(train_data, train_I_affine).to_filename("train_mask_test.nii.gz")
    # nibabel.Nifti1Image(label_data, label_I_affine).to_filename("label_mask_test.nii.gz")

    # debug = np.unique(label_data)
    # if the length of data's shape is 3, expand the dims
    # if len(train_data.shape) == 3:
    #     train_data = np.expand_dims(train_data, axis=0)
    #     label_data = np.expand_dims(label_data, axis=0)

    # if if_normalization:
    #     train_data = train_data / (np.max(train_data) - np.min(train_data))

    my_logger.info("read data from " + data_path)
    return train_data, label_data, train_I_affine, label_I_affine, Box_boundary


def preprocess(img_list):
    if isinstance(img_list, list):
        for i in range(len(img_list)):
            if np.max(img_list[i]) - np.min(img_list[i]) != 0:
                img_list[i] = (img_list[i] - np.min(img_list[i])) / (np.max(img_list[i]) - np.min(img_list[i])) * 255
            else:
                img_list[i] = numpy.zeros_like(img_list[i])
    else:  # numpy array object
        img_list = (img_list - np.min(img_list)) / (np.max(img_list) - np.min(img_list)) * 255
        return img_list


def get_appropriate_shape(original_shape_list):
    assert (len(original_shape_list) == 3)
    denominator = 16
    new_shape_list = []
    start_list = []
    for one_shape in original_shape_list:
        while one_shape % denominator != 0:
            one_shape -= 1
        new_shape_list.append(one_shape)
    for i in range(len(original_shape_list)):
        start_list.append((original_shape_list[i] - new_shape_list[i]) // 2)
    # 前闭后开
    new_roi = [[start_list[i], start_list[i] + new_shape_list[i]] for i in range(len(start_list))]
    return new_roi


def crop_data_and_save():
    """
    This function is used to crop train data manually
    :return:
    """
    if_delete_previous = True
    patch_shape = config.ModifiedUNet3DConfigure.patch_shape
    sign = str(patch_shape[0]) + str(patch_shape[1]) + str(patch_shape[2])
    data_root_path = r"D:\dataset\vessel\Parse_2022_train_data"
    all_data_path_list = glob(os.path.join(data_root_path, "*"))
    all_data_path_list = all_data_path_list
    patch_index_list = []
    for sample_dir in all_data_path_list:
        sample_name = os.path.basename(sample_dir)
        if if_delete_previous:
            train_save_path = glob(os.path.join(sample_dir, "image", "*patch.npy"))
            label_save_path = glob(os.path.join(sample_dir, "label", "*patch.npy"))
            assert (len(train_save_path) == len(label_save_path))
            for a, b in zip(train_save_path, label_save_path):
                os.remove(a)
                os.remove(b)
                my_logger.warning("!!!Delete: " + a)
                my_logger.warning("!!!Delete: " + b)
        train_save_path = glob(os.path.join(sample_dir, "image", "*patch"))
        if len(train_save_path) > 0:
            my_logger.info(sample_name + "has been patched")
            continue
        train_img, label, train_I_affine, label_I_affine, Box_boundary = read_nii_file_with_mask(sample_dir, False)
        z_num = (Box_boundary[1] - Box_boundary[0]) // patch_shape[0]
        x_num = (Box_boundary[3] - Box_boundary[2]) // patch_shape[1]
        y_num = (Box_boundary[5] - Box_boundary[4]) // patch_shape[2]
        if z_num <= 0 or x_num <= 0 or y_num <= 0:
            my_logger.warning(sample_name + "is not big enough!")
            continue
        index = 0
        for i in range(z_num + 1):
            z_start = i * patch_shape[0] + Box_boundary[0]
            z_end = (i + 1) * patch_shape[0] + Box_boundary[0]
            if z_end > Box_boundary[1]:  # the last one
                z_end = Box_boundary[1]
                z_start = Box_boundary[1] - patch_shape[0]
            for j in range(x_num + 1):
                x_start = j * patch_shape[1] + Box_boundary[2]
                x_end = (j + 1) * patch_shape[1] + Box_boundary[2]
                if x_end > Box_boundary[3]:  # the last one
                    x_end = Box_boundary[3]
                    x_start = Box_boundary[3] - patch_shape[1]
                for k in range(y_num + 1):
                    y_start = k * patch_shape[2] + Box_boundary[4]
                    y_end = (k + 1) * patch_shape[2] + Box_boundary[4]
                    if y_end > Box_boundary[5]:  # the last one
                        y_end = Box_boundary[5]
                        y_start = Box_boundary[5] - patch_shape[2]
                    train_save_path = os.path.join(sample_dir, "image",
                                                   sample_name + "_" + str(index) + "_" + sign + "_patch")
                    my_logger.info("save train numpy array to " + train_save_path)
                    np.save(train_save_path, train_img[z_start:z_end,
                                             x_start:x_end,
                                             y_start:y_end])
                    label_save_path = os.path.join(sample_dir, "label",
                                                   sample_name + "_" + str(index) + "_" + sign + "_patch")
                    np.save(label_save_path, label[z_start:z_end,
                                             x_start:x_end,
                                             y_start:y_end])
                    my_logger.info("save label numpy array to " + label_save_path)
                    index += 1
                    patch_index_list.append(
                        [z_start, z_end, x_start, x_end, y_start, y_end])


if __name__ == '__main__':
    # read_nii_file(r"D:\dataset\vessel\Parse_2022_train_data\PA000005")
    crop_data_and_save()
